from server.rpc import ServiceDiscover, ServiceAddress, ServiceInfo
from lib.redis import Redis
from typing import Callable, Optional


class RedisDiscoverService(ServiceDiscover):
    """
    结构说明：
        s区服id_服务类型: {
            地址1：最后一次心跳时间
            地址2：最后一次心跳时间
            地址3：最后一次心跳时间
        }
    """
    def __init__(self, redis: Redis):
        self.redis = redis

    def get_all_service(self, node_name: str):
        return self.redis.hgetall(node_name)

    def register_service(self, node_name: str, service_address: str, *args, **kwargs):
        return self.redis.hset(node_name, service_address, *args)

    def remove_service(self, node_name: str, *args, **kwargs):
        return self.redis.hdel(node_name, *args)

    async def watch(self, node_name: str, callback: Optional[Callable] = None):
        pub = self.redis.pubsub()
        await pub.subscribe(node_name)
        while True:
            _, node_name, service_info = await pub.parse_response()
            if _ == "subscribe":
                # print("watch success!")
                continue

            if callback is not None:
                await callback(service_info)

    def notice(self, node_name: str, service_info: ServiceInfo):
        return self.redis.publish(node_name, service_info)
